# 使用方法
# 下载媒体文件 (图片/视频)：
# http://43.156.15.142:8005/download?url=https://jimeng.jianying.com/ai-tool/work-detail/7515774963563990312
# 或者显式指定：
# http://43.156.15.142:8005/download?type=media&url=...
# 下载提示词文本文件：
# http://43.156.15.142:8005/download?type=prompt&url=https://jimeng.jianying.com/ai-tool/work-detail/7515774963563990312
# 这样，用户就可以根据需要，直接下载原始格式的文件了。

from flask import Flask, request, send_file, jsonify
import requests
from bs4 import BeautifulSoup
import tempfile
import os
from urllib.parse import urljoin, urlparse
import re

app = Flask(__name__)

def extract_work_id_from_url(target_url):
    """
    从即梦 AI URL 中提取作品 ID。
    例如: 从 'https://jimeng.jianying.com/ai-tool/work-detail/7515774963563990312?...' 提取 '7515774963563990312'
    """
    match = re.search(r'/work-detail/(\d+)', target_url)
    if match:
        return match.group(1)
    else:
        # 如果正则匹配失败，使用 URL 的哈希值作为后备方案
        import hashlib
        return hashlib.md5(target_url.encode()).hexdigest()[:10]

def fetch_content_from_url(target_url):
    """
    从给定的即梦 AI URL 中提取媒体文件（图片或视频）和提示词。
    返回: (媒体文件二进制内容, 媒体文件扩展名, 提示词文本, 作品ID)
    """
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
    }
    # 获取网页内容
    response = requests.get(target_url, headers=headers, timeout=30)
    response.raise_for_status()
    # 解析 HTML
    soup = BeautifulSoup(response.text, 'html.parser')
    # 提取作品ID用于文件命名
    work_id = extract_work_id_from_url(target_url)
    # 1. 查找提示词
    prompt_element = soup.find('span', class_='prompt-value-container-HNplhY')
    # prompt_text = prompt_element.get_text(strip=True) if prompt_element else "未获取到提示词"
    original_text = prompt_element.get_text(strip=True) if prompt_element else "未获取到提示词"
    prompt_text = "★★中文描述★★\n"+original_text
    # 2. 优先查找 video 标签
    media_element = soup.find('video')
    media_type = 'video'
    # 3. 如果没有 video，再查找 img 标签
    if not media_element:
        media_element = soup.find('img', class_='cover-img-tJwdai') or soup.find('img', id='ssr-img')
        media_type = 'image'
    if not media_element:
        raise Exception("❌ 未在页面中找到 video 或 img 元素。")
    # 获取媒体源 URL
    media_src = media_element.get('src') or media_element.get('poster')
    if not media_src:
        raise Exception(f"❌ 找到 {media_type} 元素，但未找到有效的 'src' 或 'poster' 属性。")
    # 如果是相对路径，转换为绝对路径
    if not media_src.startswith(('http://', 'https://')):
        media_src = urljoin(target_url, media_src)
    # 下载媒体文件
    media_response = requests.get(media_src, headers=headers, timeout=30)
    media_response.raise_for_status()
    # 确定文件扩展名
    if media_type == 'video':
        default_ext = 'mp4'
    else:
        default_ext = 'jpg'
    media_extension = media_src.split('?')[0].split('.')[-1].split('#')[0] if '.' in media_src.split('?')[0] else default_ext
    media_extension = media_extension.lower()
    return media_response.content, media_extension, prompt_text, work_id

@app.route('/download', methods=['GET'])
def download_images_and_prompts():
    """
    API 接口：接收一个 'url' 参数，直接返回媒体文件（图片或视频）或提示词TXT文件。
    通过 'type' 参数指定下载类型: 'media' (默认) 或 'prompt'。
    """
    target_url = request.args.get('url')
    download_type = request.args.get('type', 'media')  # 默认下载媒体文件

    if not target_url:
        return jsonify({"error": "缺少 'url' 参数。请提供即梦 AI 的工作详情页链接。"}), 400

    # 基础的 URL 校验
    parsed = urlparse(target_url)
    if not parsed.scheme or not parsed.netloc:
        return jsonify({"error": "提供的 URL 无效。"}), 400

    try:
        # 抓取内容
        media_content, media_extension, prompt_text, work_id = fetch_content_from_url(target_url)

        # 创建临时目录
        temp_dir = tempfile.mkdtemp()

        if download_type == 'prompt':
            # 处理提示词文本文件
            txt_filename = f"{work_id}.txt"
            txt_path = os.path.join(temp_dir, txt_filename)
            with open(txt_path, 'w', encoding='utf-8') as f:
                f.write(prompt_text)

            # 发送 TXT 文件
            response = send_file(
                txt_path,
                as_attachment=True,
                download_name=txt_filename,
                mimetype='text/plain'
            )

        else:  # download_type == 'media' 或其他无效值，默认为媒体文件
            # 处理媒体文件
            media_filename = f"{work_id}.{media_extension}"
            media_path = os.path.join(temp_dir, media_filename)
            with open(media_path, 'wb') as f:
                f.write(media_content)

            # 根据文件类型设置 mimetype
            mime_types = {
                'jpg': 'image/jpeg',
                'jpeg': 'image/jpeg',
                'png': 'image/png',
                'gif': 'image/gif',
                'mp4': 'video/mp4',
                'webm': 'video/webm',
                'ogg': 'video/ogg'
            }
            mimetype = mime_types.get(media_extension, 'application/octet-stream')

            # 发送媒体文件
            response = send_file(
                media_path,
                as_attachment=True,
                download_name=media_filename,
                mimetype=mimetype
            )

        # 请求结束后删除临时文件
        @response.call_on_close
        def cleanup():
            try:
                if download_type == 'prompt':
                    os.remove(txt_path)
                else:
                    os.remove(media_path)
                os.rmdir(temp_dir)
            except Exception as e:
                print(f"清理临时文件时出错: {e}")

        return response

    except Exception as e:
        return jsonify({"error": f"处理失败: {str(e)}"}), 500

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    return jsonify({"status": "OK", "message": "Flask server is running."})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8005, debug=True)
